import cv2
import numpy as np
from argparse import ArgumentParser
import skimage


def denoise_wavelet(img, sigma):
    float_img = skimage.img_as_float(img)
    # result = skimage.restoration.denoise_wavelet(float_img, sigma=sigma)
    result = skimage.restoration.denoise_wavelet(float_img, mode='hard')
    result = skimage.img_as_ubyte(result)
    return result


def main(option):
    img = cv2.imread(option.input_file, cv2.IMREAD_GRAYSCALE)
    img = denoise_wavelet(img, option.sigma)
    cv2.imwrite(option.output_file, img)


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--input_file', type=str, default='output/noise/speckle_alpha=30.00_beta=30.00.bmp')
    parser.add_argument('--output_file', type=str, default='output/denoise/wavelet/speckle.bmp')
    parser.add_argument('--sigma', type=float, default=0.01)
    option = parser.parse_args()
    main(option)
